import pandas as pd
import numpy as np
from statsmodels.tsa.holtwinters import ExponentialSmoothing


def predict_sales(df, forecast_days=14):
    """
    预测每个SPU未来14天的销量
    参数:
        df (DataFrame): 包含历史销售数据的DataFrame
        forecast_days (int): 需要预测的天数
    返回:
        DataFrame: 包含预测结果的DataFrame
    """
    # 预处理：确保数据类型正确
    df['days_since_launch'] = df['days_since_launch'].astype(int)
    df['sale_qty'] = df['sale_qty'].astype(float)

    # 存储预测结果
    forecasts = []

    # 对每个SPU单独建模预测
    for spu in df['spu_code'].unique():
        spu_data = df[df['spu_code'] == spu].copy()

        # 按时间排序
        spu_data = spu_data.sort_values('days_since_launch')

        # 创建时间序列
        ts = spu_data.set_index('days_since_launch')['sale_qty']

        # 使用Holt-Winters模型（三重指数平滑）
        try:
            # 尝试带有季节性的模型（7天周期）
            model = ExponentialSmoothing(
                ts,
                trend='add',
                seasonal='add',
                seasonal_periods=7,
                initialization_method='estimated'
            ).fit()
        except:
            # 如果季节性建模失败，使用非季节性模型
            model = ExponentialSmoothing(
                ts,
                trend='add',
                seasonal=None,
                initialization_method='estimated'
            ).fit()

        # 预测未来销量
        forecast = model.forecast(forecast_days)

        # 收集预测结果
        for day, qty in enumerate(forecast, start=ts.index.max() + 1):
            forecasts.append({
                'spu_code': spu,
                'days_since_launch': day,
                'sale_qty_pred': max(round(qty), 0)  # 确保非负
            })

    return pd.DataFrame(forecasts)


# 使用示例
if __name__ == "__main__":
    # 从Excel读取数据
    df = pd.read_excel('enhanced_sales_data.xlsx')

    # 生成预测
    forecast_df = predict_sales(df)

    # 保存结果到Excel
    forecast_df.to_excel('sales_forecast.xlsx', index=False)

    print("预测完成! 结果已保存到 sales_forecast.xlsx")